"""
数据可视化与交互式叙事模块
支持在故事视频中插入动态图表、数据可视化和交互式元素

功能：
1. 自动从故事内容中提取数据
2. 生成多种类型的可视化（柱状图、折线图、散点图、饼图等）
3. 创建动态图表动画
4. 支持交互式叙事场景
"""

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib import font_manager
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import json
import re
from PIL import Image, ImageDraw, ImageFont
import io


class DataVisualizationAgent:
    """数据可视化代理"""
    
    def __init__(self, output_dir: Path, style: str = "default"):
        """
        Args:
            output_dir: 输出目录
            style: 可视化风格 (default, dark, colorful, minimal)
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.style = style
        
        # 设置中文字体
        self._setup_chinese_font()
        
        # 设置样式
        self._setup_style()
    
    def _setup_chinese_font(self):
        """设置中文字体支持"""
        try:
            # Windows 中文字体
            plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'SimSun']
            plt.rcParams['axes.unicode_minus'] = False
        except:
            print("⚠️  中文字体设置失败，将使用默认字体")
    
    def _setup_style(self):
        """设置可视化风格"""
        style_configs = {
            'default': {
                'figure.facecolor': 'white',
                'axes.facecolor': 'white',
                'grid.alpha': 0.3
            },
            'dark': {
                'figure.facecolor': '#1e1e1e',
                'axes.facecolor': '#2d2d2d',
                'text.color': 'white',
                'axes.labelcolor': 'white',
                'xtick.color': 'white',
                'ytick.color': 'white',
                'grid.alpha': 0.2
            },
            'colorful': {
                'figure.facecolor': '#f0f8ff',
                'axes.facecolor': 'white',
                'grid.alpha': 0.4
            },
            'minimal': {
                'figure.facecolor': 'white',
                'axes.facecolor': 'white',
                'axes.spines.top': False,
                'axes.spines.right': False,
                'grid.alpha': 0.2
            }
        }
        
        if self.style in style_configs:
            plt.rcParams.update(style_configs[self.style])
    
    def extract_data_from_story(self, pages: List[str]) -> List[Dict]:
        """
        从故事内容中提取可视化数据
        
        Returns:
            [{
                'page_index': int,
                'data_type': str (numeric/comparison/timeline/composition),
                'title': str,
                'data': dict
            }]
        """
        visualizations = []
        
        # 数字模式匹配
        number_pattern = r'(\d+(?:\.\d+)?)\s*(?:个|只|名|人|次|年|月|天|小时|分钟|秒|米|千米|公斤|元|%)'
        
        for idx, page in enumerate(pages):
            # 查找数字数据
            numbers = re.findall(number_pattern, page)
            
            if len(numbers) >= 2:
                # 提取比较数据
                viz_data = self._create_comparison_data(page, numbers, idx)
                if viz_data:
                    visualizations.append(viz_data)
            
            # 检查时间序列关键词
            if any(word in page for word in ['增长', '变化', '趋势', '发展', '历程']):
                viz_data = self._create_timeline_data(page, idx)
                if viz_data:
                    visualizations.append(viz_data)
            
            # 检查组成部分关键词
            if any(word in page for word in ['组成', '包括', '分为', '由...构成']):
                viz_data = self._create_composition_data(page, idx)
                if viz_data:
                    visualizations.append(viz_data)
        
        return visualizations
    
    def _create_comparison_data(self, text: str, numbers: List[str], page_idx: int) -> Optional[Dict]:
        """创建对比数据"""
        try:
            values = [float(n) for n in numbers[:5]]  # 最多5个数据点
            
            # 简单的标签提取（改进空间大）
            labels = [f"项目{i+1}" for i in range(len(values))]
            
            return {
                'page_index': page_idx,
                'data_type': 'comparison',
                'chart_type': 'bar',  # 柱状图
                'title': f'数据对比 - 第{page_idx+1}页',
                'data': {
                    'labels': labels,
                    'values': values
                }
            }
        except:
            return None
    
    def _create_timeline_data(self, text: str, page_idx: int) -> Optional[Dict]:
        """创建时间序列数据"""
        # 这里简化处理，实际可以用NLP提取
        return {
            'page_index': page_idx,
            'data_type': 'timeline',
            'chart_type': 'line',
            'title': f'发展趋势 - 第{page_idx+1}页',
            'data': {
                'labels': ['阶段1', '阶段2', '阶段3', '阶段4'],
                'values': [10, 25, 45, 80]  # 示例数据
            }
        }
    
    def _create_composition_data(self, text: str, page_idx: int) -> Optional[Dict]:
        """创建组成数据"""
        return {
            'page_index': page_idx,
            'data_type': 'composition',
            'chart_type': 'pie',
            'title': f'组成分析 - 第{page_idx+1}页',
            'data': {
                'labels': ['部分A', '部分B', '部分C', '其他'],
                'values': [30, 25, 20, 25]
            }
        }
    
    def create_bar_chart(self, data: Dict, output_path: Path, animated: bool = False) -> str:
        """
        创建柱状图
        
        Args:
            data: 数据字典
            output_path: 输出路径
            animated: 是否创建动画版本
        """
        labels = data['data']['labels']
        values = data['data']['values']
        
        fig, ax = plt.subplots(figsize=(12, 7))
        
        if animated:
            # 创建动画
            bars = ax.bar(labels, [0] * len(values), color='#3498db', alpha=0.8)
            ax.set_ylim(0, max(values) * 1.2)
            ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
            ax.set_xlabel('类别', fontsize=14)
            ax.set_ylabel('数值', fontsize=14)
            ax.grid(axis='y', alpha=0.3, linestyle='--')
            
            def animate(frame):
                progress = frame / 30  # 30帧动画
                for bar, value in zip(bars, values):
                    bar.set_height(value * progress)
                return bars
            
            anim = animation.FuncAnimation(fig, animate, frames=31, interval=50, blit=True)
            anim.save(str(output_path).replace('.png', '.gif'), writer='pillow', fps=20)
            
        else:
            # 静态图
            bars = ax.bar(labels, values, color='#3498db', alpha=0.8, edgecolor='navy')
            
            # 添加数值标签
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height,
                       f'{height:.1f}',
                       ha='center', va='bottom', fontsize=12, fontweight='bold')
            
            ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
            ax.set_xlabel('类别', fontsize=14)
            ax.set_ylabel('数值', fontsize=14)
            ax.grid(axis='y', alpha=0.3, linestyle='--')
        
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        return str(output_path)
    
    def create_line_chart(self, data: Dict, output_path: Path, animated: bool = False) -> str:
        """创建折线图"""
        labels = data['data']['labels']
        values = data['data']['values']
        
        fig, ax = plt.subplots(figsize=(12, 7))
        
        if animated:
            # 动画折线图
            line, = ax.plot([], [], marker='o', linewidth=3, 
                           markersize=10, color='#e74c3c', alpha=0.8)
            ax.set_xlim(-0.5, len(labels) - 0.5)
            ax.set_ylim(0, max(values) * 1.2)
            ax.set_xticks(range(len(labels)))
            ax.set_xticklabels(labels)
            ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
            ax.grid(True, alpha=0.3, linestyle='--')
            
            def animate(frame):
                progress = (frame + 1) / 30
                visible_points = int(len(labels) * progress)
                x = range(visible_points)
                y = values[:visible_points]
                line.set_data(x, y)
                return line,
            
            anim = animation.FuncAnimation(fig, animate, frames=30, interval=100, blit=True)
            anim.save(str(output_path).replace('.png', '.gif'), writer='pillow', fps=10)
            
        else:
            # 静态折线图
            ax.plot(labels, values, marker='o', linewidth=3, 
                   markersize=10, color='#e74c3c', alpha=0.8)
            
            # 添加数值标签
            for i, (label, value) in enumerate(zip(labels, values)):
                ax.text(i, value, f'{value:.1f}', 
                       ha='center', va='bottom', fontsize=11, fontweight='bold')
            
            ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
            ax.grid(True, alpha=0.3, linestyle='--')
        
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        return str(output_path)
    
    def create_pie_chart(self, data: Dict, output_path: Path, animated: bool = False) -> str:
        """创建饼图"""
        labels = data['data']['labels']
        values = data['data']['values']
        
        fig, ax = plt.subplots(figsize=(10, 10))
        
        colors = ['#3498db', '#e74c3c', '#2ecc71', '#f39c12', '#9b59b6']
        explode = [0.05] * len(values)  # 分离效果
        
        if animated:
            # 动画饼图（从0度旋转到360度）
            def animate(frame):
                ax.clear()
                start_angle = 90 - (frame * 12)  # 每帧旋转12度
                ax.pie(values, labels=labels, autopct='%1.1f%%',
                      startangle=start_angle, colors=colors, explode=explode,
                      textprops={'fontsize': 12, 'fontweight': 'bold'})
                ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
            
            anim = animation.FuncAnimation(fig, animate, frames=30, interval=50)
            anim.save(str(output_path).replace('.png', '.gif'), writer='pillow', fps=20)
            
        else:
            # 静态饼图
            ax.pie(values, labels=labels, autopct='%1.1f%%',
                  startangle=90, colors=colors, explode=explode,
                  textprops={'fontsize': 12, 'fontweight': 'bold'},
                  shadow=True)
            ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
        
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        return str(output_path)
    
    def create_scatter_plot(self, data: Dict, output_path: Path) -> str:
        """创建散点图"""
        # 生成示例数据
        np.random.seed(42)
        x = np.random.randn(50) * 10 + 50
        y = x * 0.8 + np.random.randn(50) * 5
        
        fig, ax = plt.subplots(figsize=(12, 8))
        
        scatter = ax.scatter(x, y, s=100, c=y, cmap='viridis', 
                            alpha=0.6, edgecolors='black', linewidth=1)
        
        # 添加趋势线
        z = np.polyfit(x, y, 1)
        p = np.poly1d(z)
        ax.plot(x, p(x), "r--", alpha=0.8, linewidth=2, label='趋势线')
        
        ax.set_title(data['title'], fontsize=18, fontweight='bold', pad=20)
        ax.set_xlabel('X轴', fontsize=14)
        ax.set_ylabel('Y轴', fontsize=14)
        ax.legend(fontsize=12)
        ax.grid(True, alpha=0.3, linestyle='--')
        
        plt.colorbar(scatter, label='数值')
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        return str(output_path)
    
    def create_interactive_infographic(self, data: Dict, output_path: Path) -> str:
        """创建交互式信息图"""
        # 使用PIL创建信息图
        width, height = 1920, 1080
        img = Image.new('RGB', (width, height), color='#f0f0f0')
        draw = ImageDraw.Draw(img)
        
        try:
            # 尝试加载中文字体
            title_font = ImageFont.truetype("msyh.ttc", 60)
            text_font = ImageFont.truetype("msyh.ttc", 36)
        except:
            title_font = ImageFont.load_default()
            text_font = ImageFont.load_default()
        
        # 绘制标题
        title = data.get('title', '信息图')
        title_bbox = draw.textbbox((0, 0), title, font=title_font)
        title_width = title_bbox[2] - title_bbox[0]
        draw.text(((width - title_width) // 2, 50), title, 
                 fill='#2c3e50', font=title_font)
        
        # 绘制数据卡片
        labels = data['data']['labels']
        values = data['data']['values']
        
        card_width = 350
        card_height = 200
        cards_per_row = 4
        margin = 50
        
        for i, (label, value) in enumerate(zip(labels, values)):
            row = i // cards_per_row
            col = i % cards_per_row
            
            x = margin + col * (card_width + margin)
            y = 200 + row * (card_height + margin)
            
            # 绘制卡片
            draw.rectangle([x, y, x + card_width, y + card_height],
                          fill='white', outline='#3498db', width=3)
            
            # 绘制标签
            label_bbox = draw.textbbox((0, 0), label, font=text_font)
            label_width = label_bbox[2] - label_bbox[0]
            draw.text((x + (card_width - label_width) // 2, y + 30),
                     label, fill='#2c3e50', font=text_font)
            
            # 绘制数值
            value_text = f"{value:.1f}"
            value_bbox = draw.textbbox((0, 0), value_text, font=title_font)
            value_width = value_bbox[2] - value_bbox[0]
            draw.text((x + (card_width - value_width) // 2, y + 100),
                     value_text, fill='#e74c3c', font=title_font)
        
        img.save(output_path, quality=95)
        return str(output_path)
    
    def generate_all_visualizations(self, pages: List[str], animated: bool = True) -> List[str]:
        """
        为故事生成所有可视化
        
        Args:
            pages: 故事页面列表
            animated: 是否生成动画版本
        
        Returns:
            生成的可视化图片路径列表
        """
        print("\n" + "="*70)
        print("🎨 数据可视化生成")
        print("="*70)
        
        # 提取数据
        viz_data_list = self.extract_data_from_story(pages)
        
        if not viz_data_list:
            print("ℹ️  未检测到可视化数据，生成示例图表...")
            viz_data_list = self._create_demo_visualizations(len(pages))
        
        print(f"📊 检测到 {len(viz_data_list)} 个可视化机会")
        
        generated_paths = []
        
        for idx, viz_data in enumerate(viz_data_list):
            chart_type = viz_data.get('chart_type', 'bar')
            output_name = f"viz_{idx+1:03d}_{chart_type}.png"
            output_path = self.output_dir / output_name
            
            print(f"\n[{idx+1}/{len(viz_data_list)}] 生成 {chart_type} 图表...")
            
            try:
                if chart_type == 'bar':
                    path = self.create_bar_chart(viz_data, output_path, animated)
                elif chart_type == 'line':
                    path = self.create_line_chart(viz_data, output_path, animated)
                elif chart_type == 'pie':
                    path = self.create_pie_chart(viz_data, output_path, animated)
                elif chart_type == 'scatter':
                    path = self.create_scatter_plot(viz_data, output_path)
                else:
                    path = self.create_interactive_infographic(viz_data, output_path)
                
                generated_paths.append(path)
                print(f"✅ 已保存: {output_path.name}")
                
            except Exception as e:
                print(f"❌ 生成失败: {e}")
        
        print(f"\n{'='*70}")
        print(f"✅ 可视化生成完成！共 {len(generated_paths)} 个图表")
        print(f"📁 保存位置: {self.output_dir}")
        print(f"{'='*70}\n")
        
        return generated_paths
    
    def _create_demo_visualizations(self, num_pages: int) -> List[Dict]:
        """创建演示可视化数据"""
        demos = [
            {
                'page_index': 0,
                'chart_type': 'bar',
                'title': '角色能力对比',
                'data': {
                    'labels': ['力量', '智慧', '速度', '魔法', '防御'],
                    'values': [85, 92, 78, 95, 88]
                }
            },
            {
                'page_index': num_pages // 2,
                'chart_type': 'line',
                'title': '冒险进度',
                'data': {
                    'labels': ['起点', '第一关', '第二关', '第三关', '终点'],
                    'values': [0, 25, 50, 75, 100]
                }
            },
            {
                'page_index': num_pages - 1,
                'chart_type': 'pie',
                'title': '资源分配',
                'data': {
                    'labels': ['金币', '装备', '药水', '技能书'],
                    'values': [40, 30, 20, 10]
                }
            }
        ]
        
        return demos[:min(3, num_pages)]
    
    def extract_data_from_story(self, pages: List[str]) -> List[Dict]:
        """自定义数据提取逻辑"""
        visualizations = []
    
        for idx, page in enumerate(pages):
         # 示例：检测特定关键词
            if "对比" in page or "比较" in page:
            # 提取对比数据
              viz_data = self._extract_comparison(page, idx)
              visualizations.append(viz_data)
        
            if "增长" in page or "变化" in page:
            # 提取趋势数据
               viz_data = self._extract_trend(page, idx)
               visualizations.append(viz_data)
    
        return visualizations
    def _extract_comparison(self, text: str, page_idx: int) -> Optional[Dict]:
            """提取对比数据 - 实现版本"""
            # 尝试提取数字
            number_pattern = r'(\d+(?:\.\d+)?)\s*(?:个|只|名|人|次|年|月|天|小时|分钟|秒|米|千米|公斤|元|%|分|点)'
            numbers = re.findall(number_pattern, text)
            
            if len(numbers) >= 2:
                try:
                    values = [float(n) for n in numbers[:5]]
                    labels = [f"项目{i+1}" for i in range(len(values))]
                    
                    return {
                        'page_index': page_idx,
                        'data_type': 'comparison',
                        'chart_type': 'bar',
                        'title': f'数据对比 - 第{page_idx+1}页',
                        'data': {
                            'labels': labels,
                            'values': values
                        }
                    }
                except:
                    pass
            
            return None
        
    def _extract_trend(self, text: str, page_idx: int) -> Optional[Dict]:
        """提取趋势数据 - 实现版本"""
        # 尝试提取数字序列
        number_pattern = r'(\d+(?:\.\d+)?)'
        numbers = re.findall(number_pattern, text)
        
        if len(numbers) >= 3:
            try:
                values = [float(n) for n in numbers[:6]]
                labels = [f"阶段{i+1}" for i in range(len(values))]
                
                return {
                    'page_index': page_idx,
                    'data_type': 'timeline',
                    'chart_type': 'line',
                    'title': f'发展趋势 - 第{page_idx+1}页',
                    'data': {
                        'labels': labels,
                        'values': values
                    }
                }
            except:
                pass
        
        return None

# ==================== 使用示例 ====================

def integrate_with_story_agent(story_dir: Path, pages: List[str], 
                               insert_viz: bool = True) -> List[str]:
    """
    集成到故事生成器的辅助函数
    
    Args:
        story_dir: 故事目录
        pages: 故事页面
        insert_viz: 是否插入可视化
    
    Returns:
        可视化图片路径列表
    """
    if not insert_viz:
        return []
    
    viz_agent = DataVisualizationAgent(
        output_dir=story_dir / "visualizations",
        style="colorful"
    )
    
    return viz_agent.generate_all_visualizations(pages, animated=True)


if __name__ == "__main__":
    # 测试示例
    test_pages = [
        "小明的数学成绩从60分提高到了95分，他的努力得到了回报。",
        "班级里有30%的同学喜欢数学，25%喜欢语文，20%喜欢英语，25%喜欢其他科目。",
        "销售额在过去四个季度分别为100万、150万、200万和280万元，呈现持续增长趋势。"
    ]
    
    output_dir = Path("test_visualizations")
    viz_agent = DataVisualizationAgent(output_dir, style="colorful")
    
    paths = viz_agent.generate_all_visualizations(test_pages, animated=True)
    print(f"\n生成的图表: {paths}")